# -*- coding: utf-8 -*-
import re

# 关系字典配置
voice_mapping = {
    "师傅": "云健(成年男)",
    "老者": "云野(老年男)",
    "小孩": "云夏(儿童男)",
    "大徒弟": "云健(成年男)",
    "二徒弟": "云皓(成年男)",
    "强盗": "云健(成年男)",
    "师傅的女儿": "晓伊(年轻女)",
    "助手": "晓晓(年轻女)",
    "女儿": "晓伊(年轻女)",
    "爸爸": "云健(成年男)",
    "小智": "云健(成年男)",
    "皮卡丘": "晓伊(年轻女)",
   
    # 可以继续添加其他角色映射
}

# 每个角色的情感选项
emotion_options = {
    "云希(年轻男)": ["narration-relaxed", "embarrassed", "fearful", "cheerful", "disgruntled", "serious", "angry", "sad", "depressed", "chat", "assistant", "newscast"],
    "云健(成年男)": ["narration-relaxed", "sports-commentary", "sports-commentary-excited", "angry", "disgruntled", "cheerful", "sad", "serious", "depressed", "documentary-narration"],
    "云扬(成年男)": ["customerservice", "narration-professional", "newscast-casual"],
    "云皓(成年男)": ["advertisement-upbeat"],
    "云夏(儿童男)": ["calm", "fearful", "cheerful", "angry", "sad"],
    "云野(老年男)": ["embarrassed", "calm", "fearful", "cheerful", "disgruntled", "serious", "angry", "sad"],
    "云泽(老年男)": ["calm", "fearful", "cheerful", "disgruntled", "serious", "angry", "sad", "depressed", "documentary-narration"],
    "晓晓(年轻女)": ["assistant", "chat", "customerservice", "newscast", "affectionate", "angry", "calm", "cheerful", "disgruntled", "fearful", "gentle", "lyrical", "sad", "serious", "poetry-reading", "friendly", "chat-casual", "whispering", "sorry", "excited"],
    "晓伊(年轻女)": ["angry", "disgruntled", "affectionate", "cheerful", "fearful", "sad", "embarrassed", "serious", "gentle"],
    "晓辰(年轻女)": ["livecommercial"],
    "晓涵(成年女)": ["calm", "fearful", "cheerful", "disgruntled", "serious", "angry", "sad", "gentle", "affectionate", "embarrassed"],
    "晓梦(年轻女)": ["chat"],
    "晓墨(成年女)": ["embarrassed", "calm", "fearful", "cheerful", "disgruntled", "serious", "angry", "sad", "depressed", "affectionate", "gentle", "envious"],
    "晓睿(老年女)": ["calm", "fearful", "angry", "sad"],
    "晓双(儿童女)": ["chat"],
    "晓甄(年轻女)": ["angry", "disgruntled", "cheerful", "fearful", "sad", "serious"],
}

# 强度选项
strength_options = {
    "强": "2",
    "中": "1",
    "弱": "0.5",
    "强烈": "2",
    "中等": "1",
    "轻微": "0.5",
}

# 情感描述映射
emotion_mapping = {
    "兴奋": "cheerful",
    "中等": "unknown",
    "轻微": "unknown",
    "强烈": "unknown",
    "愤怒": "angry",
    "担忧": "fearful",
    "惊讶": "surprised",
    "疑惑": "confused",
    "微笑": "happy",
    "威胁": "threatening",
    "请求": "requesting",
    "抱怨": "complaining",
    "解释": "explaining",
    "安慰": "comforting",
    "命令": "commanding",
    "发现": "discovering",
    "同情": "sympathetic",
}

def get_emotion_for_character(voice):
    """根据角色获取合适的情感"""
    print(f"voice: {voice}")

    if voice in emotion_options:
        emotions = emotion_options[voice]
        return emotions[0] if emotions else "calm"  # 返回第一个情感选项，如果没有则返回unknown
    return "unknown"

def extract_dialogues(file_path):
    """从 OCR 结果文件中提取对话并格式化为指定结构"""
    dialogues = []
    current_line_number = 0
    i = 0
    
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            current_line_number += 1
            line = line.strip()
            
            # 跳过空行和分隔线
            if not line or line.startswith('='):
                continue
                
            # 跳过图片标题行
            if line.startswith('图片:'):
                continue
            
            # 使用正则表达式匹配对话格式（使用@分隔）
            parts = line.split('@')
            if len(parts) >= 3:  # 至少需要编号、角色和对话内容
                try:
                    number = parts[0].strip()
                    character = parts[1].strip()
                    dialogue = parts[-1].strip()  # 取最后一部分作为对话内容
                    i += 1
                    
                    # 获取语音、感情和强度
                    voice = voice_mapping.get(character, "云健(成年男)")
                    emotion_description = parts[2].strip()  # 获取情感描述
                    strength_description = parts[3].strip()  # 获取强度描述
                    emotion = emotion_mapping.get(emotion_description, "calm")  # 获取情感
                    #如果该角色没有这个感情，则使用该角色第一个感情
                    if emotion not in emotion_options[voice]:
                        emotion = emotion_options[voice][0]

                    strength = strength_options.get(strength_description, "default")  # 获取强度
                    
                    # 格式化对话
                    formatted_dialogue = f"台词{i}@{character}@语音: {voice}@感情: {emotion}@强度: {strength}@{dialogue}"
                    dialogues.append(formatted_dialogue)
                except Exception as e:
                    print(f"警告：第 {current_line_number} 行处理出错: {line}")
                    print(f"错误信息: {str(e)}")
            else:
                print(f"警告：第 {current_line_number} 行格式不正确: {line}")
    
    return dialogues

if __name__ == "__main__":
    # 打印原始内容
    # print("原始文件内容：")
    # with open("ocr_results.txt", 'r', encoding='utf-8') as f:
    #     print(f.read())
    
    print("\n处理后的对话：")
    dialogues = extract_dialogues("ocr_results.txt")
    # 将对话内容保存到txt文件中
    with open("/mymount/xuniji/mypacong/data/dialogues.txt", "w", encoding="utf-8") as f:
        for dialogue in dialogues:
            f.write(dialogue + "\n")
    with open("./dialogues.txt", "w", encoding="utf-8") as f:
        for dialogue in dialogues:
            f.write(dialogue + "\n")
    print("test_dialogues = [")
    for dialogue in dialogues:
        print(f'    "{dialogue}",')
    print("]") 
    print_dialogues = "test_dialogues = [" + "\n".join("'" + dialogue + "'" for dialogue in dialogues) + "]"
    # 将print输出的内容保存到taichi.py文件中,如果文件存在则覆盖
    with open("taichi.py", "w", encoding="utf-8") as f:
        f.write(print_dialogues)
